import torch
import torch.nn as nn
from feature_extract.registry import MODEL_REGISTRY
from feature_extract.model.light_reid import resnet50ibna
from feature_extract.model.light_reid.skeleton_head import SkeletonMultiHeads

__all__ = ['SkeletonReIDModel']


@MODEL_REGISTRY.register('SkeletonReIDModel')
class SkeletonReIDModel(nn.Module):
    """
    Architecture for ReID Model (using skeleton)
        combine backbone, pooling and head modules
        remove pooling, i.e. img --> backbone --> feature map --> head
    """

    def __init__(self, class_num=1, pretrained=False, device=torch.device('cpu'), **kwargs):
        super(SkeletonReIDModel, self).__init__()
        self.backbone = resnet50ibna(pretrained=pretrained, last_stride_one=True)
        self.pooling = nn.AdaptiveAvgPool2d(1)
        self.head = SkeletonMultiHeads(class_num=class_num, in_dim=self.backbone.dim, device=device, **kwargs)
        self.disable_hash()

    def forward(self, x, y=None, fixcnn=False):
        '''
        Args:
            x(torch.tensor): images
            y(torch.tensor): labels, required by circle_softmax, arc_softmax
            fixcnn(bool): if True, detach feature map
        '''
        # cnn backbone
        feats_map = self.backbone(x)
        if fixcnn:
            feats_map = feats_map.detach()
        ## pooling
        # feats_vec = self.pooling(feats_map).squeeze(3).squeeze(2) # support bs=1
        # head
        res = self.head(x, feats_map, y)

        # return
        res['feats_map'] = feats_map
        # res['feats_vec'] = feats_vec

        if self.training:
            return res
        else:
            return res[self.test_feats]

    def enable_hash(self):
        self.head.enable_hash()
        self.test_feats = 'binary_feats'

    def disable_hash(self):
        self.head.disable_hash()
        self.test_feats = 'bn_feats'
